#include "server/websocket_server.hpp"

bool websocket_connection::_send_handshake_response(const std::string& key) const {
    handshake_response response;
    response.version = "HTTP/1.1";
    response.status = "101";
    response.reason = "Switching Protocols";
    response.message.key_values["Upgrade"] = {"websocket"};
    response.message.key_values["Connection"] = {"Upgrade"};
    response.message.key_values["Sec-WebSocket-Accept"] = {sha1_then_base64(key + WEBSOCKET_SERVER_MAGIC_KEY)};
    auto str = response.dump();
    if (!server_instance._server.send(sock, std::vector(str.data(), str.data() + str.length()))) {
        lerror("Sending handshake response failed.");
        return false;
    }
    return true;
}

bool websocket_connection::_send_pong(std::vector<char> buf) const {
    websocket_frame frame;
    frame.opcode = PONG;
    frame.length = buf.size();
    frame.buffer = std::move(buf);
    return server_instance._server.send(sock, frame.dump());
}

bool websocket_connection::_handle_websocket_frame(websocket_frame frame) {
    _last_recv_time = std::chrono::time_point_cast<std::chrono::milliseconds>(std::chrono::system_clock::now());
    switch (frame.opcode) {
    case TEXT_FRAME: {
        linfo("received text frame from client " << nonstd::to_string(address) << " with size " << frame.length);
        server_instance._recv_hook(sock, address, std::move(frame));
        break;
    }
    case BINARY_FRAME: {
        linfo("received binary frame from client " << nonstd::to_string(address) << " with size " << frame.length);
        server_instance._recv_hook(sock, address, std::move(frame));
        break;
    }
    case PING: {
        linfo("received ping from client " << nonstd::to_string(address));
        auto buf = frame.buffer;
        server_instance._recv_hook(sock, address, std::move(frame));
        return _send_pong(std::move(buf));
    }
    case PONG: {
        linfo("received pong from client " << nonstd::to_string(address));
        server_instance._recv_hook(sock, address, std::move(frame));
        break;
    }
    case CONNECTION_CLOSE: {
        if (state == websocket_state::CLOSING) {
            linfo("received websocket close confirm from client " << nonstd::to_string(address));
            state = websocket_state::CLOSED;
        } else {
            linfo("received websocket close request from client " << nonstd::to_string(address));
            auto buf = frame.buffer;
            server_instance._recv_hook(sock, address, std::move(frame));
            server_instance.send(sock, std::move(buf), CONNECTION_CLOSE, true, [&, this]() {
                ldebug("websocket close confirm sent to client.");
                state = websocket_state::CLOSED;
                server_instance._server.close_socket(this->sock);
            });
            state = websocket_state::CLOSING;
            return true;
        }
        return false;
    }
    default:
        lerror("unexpected websocket frame with opcode: " << frame.opcode << " from " << nonstd::to_string(address));
        return false;
    }
    return true;
}

websocket_connection::~websocket_connection() = default;

bool websocket_connection::_recv(std::vector<char> buf) {
    if (buf.empty() || state == websocket_state::CLOSED) return true;
    ldebug("received " << buf.size() << " bytes from client " << nonstd::to_string(address));
    if (state == websocket_state::CONNECTING) {
        _handshake_buffer += std::string(buf.data(), buf.size());
        if (_handshake_buffer.find("\r\n") == std::string::npos) {
            return true;
        }
        auto handshake_result = handshake_request::parse_handshake_request(_handshake_buffer);
        if (handshake_result.index() == 0) {
            lerror("invalid handshake message from " << nonstd::to_string(address));
            lerror("handshake error: " << std::get<std::string>(handshake_result));
            return false;
        }
        auto header = std::get<handshake_request>(handshake_result);
        auto key = header.message.key_values["Sec-WebSocket-Key"][0];
        if (!_send_handshake_response(key)) {
            return false;
        }
        state = websocket_state::OPEN;
        linfo("accepted websocket handshake from " << nonstd::to_string(address)
                                                   << ", websocket connection established.");
        server_instance._accept_hook(*const_cast<websocket_connection*>(this));
        return true;
    }
    return _recv_buffer.receive_data(std::move(buf));
}

void websocket_server::start() {

    _server.start();
}

bool websocket_server::send(win::socket sock, std::vector<char> buf, websocket_frame_code opcode, bool FIN,
    const std::function<void()>& callback) {
    if (opcode > 7 && !FIN) {
        lwarn("control frame must be finished.");
        return false;
    }
    auto connection = _get_connection(sock);
    if (!connection) {
        lwarn("cannot find connection with socket " << sock);
        return false;
    }
    if (!connection->available()) {
        lwarn("websocket connection is not available with socket " << sock);
        return false;
    }
    websocket_frame frame;
    frame.opcode = opcode;
    frame.FIN = FIN;
    frame.length = buf.size();
    frame.buffer = std::move(buf);
    auto dumped = frame.dump();
    return _server.send(sock, std::move(dumped), callback);
}

bool websocket_server::close(win::socket sock, int close_cause) {
    auto conn = _get_connection(sock);
    if (!conn) {
        lwarn("cannot close socket " << sock << ", cannot find such socket");
        return false;
    }
    websocket_frame frame;
    frame.opcode = CONNECTION_CLOSE;
    frame.masked = true;
    frame.length = 2;
    frame.buffer = {static_cast<char>(close_cause / 256), static_cast<char>(close_cause % 256)};
    conn->state = websocket_state::CLOSING;
    return _server.send(sock, frame.dump());
}

bool websocket_server::stop() {
    using namespace std::literals;
    _heartbeat_exit = true;
    _server.stop();
    linfo("websocket server stopped.");
    return true;
}

websocket_connection& websocket_server::_add_new_connection(win::socket sock, const win::sockaddr_in& address) {
    std::scoped_lock<std::mutex> lock(_m);
    _connections.emplace(sock, std::make_unique<websocket_connection>(*this, sock, address));
    return *_connections[sock];
}

websocket_connection* websocket_server::_get_connection(win::socket sock) {
    std::scoped_lock<std::mutex> lock(_m);
    auto res = _connections.find(sock);
    if (res == _connections.end()) {
        return nullptr;
    }
    return res->second.get();
}

bool websocket_server::_remove_connection(win::socket sock) {
    std::scoped_lock<std::mutex> lock(_m);
    auto res = _connections.find(sock);
    if (res == _connections.end()) {
        return false;
    }
    auto addr = res->second->address;
    _connections.erase(res);
    if (_close_hook) {
        _close_hook(sock, addr);
    }
    return true;
}

void websocket_server::_initialize_socket_server() {
    _recv_hook = [](win::socket sock, const win::sockaddr_in& addr, websocket_frame frame) {
        linfo("received websocket data from socket " << sock << " : " << nonstd::to_string(addr) << "\n"
                                                     << frame.description());
    };
    _accept_hook = [](websocket_connection& conn) {
        linfo("accepted websocket handshake from socket " << conn.sock << " : " << nonstd::to_string(conn.address));
    };

    _server.register_accept_hooks(
        [this](win::socket sock, const win::sockaddr_in& addr) { _add_new_connection(sock, addr); });

    _server.register_recv_hooks([this](win::socket sock, const win::sockaddr_in& addr, std::vector<char> buf) -> bool {
        auto ws_connection = _get_connection(sock);
        while (!ws_connection) {
            using namespace std::literals;
            std::this_thread::sleep_for(100us);
            ws_connection = _get_connection(sock);
        }
        if (!ws_connection->_recv(std::move(buf))) {
            ws_connection->state = websocket_state::CLOSING;
            auto close_res = _server.close_socket(sock);
            assert(close_res, "cannot remove socket connection with address "
                                  << nonstd::to_string(addr) << " after websocket receive failure or websocket close.");
            return false;
        }
        // ldebug("websocket server handle recevied data success.");
        return true;
    });

    _server.register_close_hooks([this](win::socket sock, const win::sockaddr_in& addr) {
        if (_remove_connection(sock)) {
            linfo("websocket connection with socket " << sock << " from " << nonstd::to_string(addr) << " closed");
        }
    });

    if constexpr (ENABLE_HEARTBEAT) {
        _heartbeat_exit = false;
        _heartbeat_thread = std::thread([this]() {
            using namespace std::literals;
            while (!_heartbeat_exit) {
                std::unique_lock<std::mutex> uni_lock(_heartbeat_mutex);
                _heartbeat_cv.wait_for(uni_lock, HEARTBEAT_INTERVAL, [&] { return _heartbeat_exit.load(); });
                if (_heartbeat_exit) {
                    break;
                }
                auto now = std::chrono::time_point_cast<std::chrono::milliseconds>(std::chrono::system_clock::now());
                std::vector<win::socket> unresponse_sockets;
                {
                    std::scoped_lock<std::mutex> lock(_m);
                    for (const auto& conn: _connections) {
                        auto dur = conn.second->_last_recv_time - now;
                        if (dur > HEARTBEAT_INTERVAL * 4) {
                            unresponse_sockets.push_back(conn.second->sock);
                            linfo("websocket connection " << conn.second->sock << " timeout, closing.");
                        } else if (dur > HEARTBEAT_INTERVAL * 2) {
                            websocket_frame frame;
                            frame.opcode = PING;
                            if (!_server.send(conn.second->sock, frame.dump())) {
                                lerror("send ping to client failed");
                            }
                        }
                    }
                }
                for (const auto& sock: unresponse_sockets) {
                    _server.close_socket(sock);
                }
            }
            linfo("websocket heartbeat thread closed.");
        });
    }
}
